#!/usr/bin/env python3
"""
AI-Powered SSH Terminal - Python Backend
Enterprise-grade terminal with local AI model support
"""

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, FileResponse
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
import asyncio
import json
import logging
from pathlib import Path
from datetime import datetime

# Import our modules
from ai_engine import AICommandGenerator, AIConnectionManager
from ssh_manager import SSHConnectionManager
from database import DatabaseManager
from models import *

# Configure logging with more detailed format
logging.basicConfig(
    level=logging.DEBUG,  # More verbose logging
    format='%(asctime)s - %(name)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s',
    handlers=[
        logging.StreamHandler(),  # Console output
        logging.FileHandler('ai_terminal.log', encoding='utf-8')  # File output
    ]
)
logger = logging.getLogger(__name__)

# Initialize FastAPI app
app = FastAPI(
    title="AI-Powered SSH Terminal",
    description="Enterprise SSH terminal with local AI assistance",
    version="2.0.0"
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global managers
db_manager = DatabaseManager()
ssh_manager = SSHConnectionManager()
ai_generator = AICommandGenerator()
ai_connection_manager = AIConnectionManager()

# Active WebSocket connections
active_connections: dict = {}

@app.on_event("startup")
async def startup_event():
    """Initialize application on startup"""
    logger.info("Starting AI-Powered SSH Terminal...")
    
    # Initialize database
    db_manager.init_database()
    
    # AI engine already initializes Ollama automatically in __init__
    logger.info(f"AI Engine Status: {ai_generator.service_type} - Model loaded: {ai_generator.model_loaded}")
    
    logger.info("Application started successfully")

@app.on_event("shutdown")
async def shutdown_event():
    """Cleanup on shutdown"""
    logger.info("Shutting down application...")
    
    # Close all SSH connections
    ssh_manager.close_all_connections()
    
    # Close database connections
    db_manager.close()

# Serve static files
app.mount("/static", StaticFiles(directory="static"), name="static")

@app.get("/")
async def read_root():
    """Serve main application"""
    return FileResponse("ai-terminal.html")

@app.get("/favicon.ico")
async def favicon():
    """Serve favicon"""
    return FileResponse("static/favicon.ico")

@app.get("/debug")
async def debug_info():
    """Debug endpoint to check system status"""
    logger.info("Debug endpoint accessed")
    debug_data = {
        "status": "healthy",
        "ai_model_loaded": ai_generator.model_loaded,
        "active_websocket_connections": len(active_connections),
        "active_ssh_sessions": len(ssh_manager.sessions),
        "ssh_sessions_detail": {
            session_id: {
                "hostname": session.connection_info.hostname,
                "username": session.connection_info.username,
                "connected": session.connected
            } for session_id, session in ssh_manager.sessions.items()
        },
        "websocket_clients": list(active_connections.keys())
    }
    logger.debug(f"Debug response: {debug_data}")
    return debug_data

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    logger.info("Health check requested")
    health_data = {
        "status": "healthy",
        "ai_model_loaded": ai_generator.model_loaded,
        "active_connections": len(active_connections),
        "ssh_sessions": len(ssh_manager.sessions)
    }
    logger.debug(f"Health check response: {health_data}")
    return health_data

@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    """WebSocket endpoint for terminal communication"""
    logger.info(f"New WebSocket connection attempt from client: {client_id}")
    await websocket.accept()
    active_connections[client_id] = websocket
    logger.info(f"WebSocket connection established for client: {client_id}. Total connections: {len(active_connections)}")
    
    try:
        while True:
            data = await websocket.receive_json()
            logger.debug(f"Received message from {client_id}: {data.get('type', 'unknown')}")
            await handle_websocket_message(websocket, client_id, data)
            
    except WebSocketDisconnect:
        logger.info(f"Client {client_id} disconnected normally")
        if client_id in active_connections:
            del active_connections[client_id]
        ssh_manager.cleanup_client_connections(client_id)
        logger.info(f"Cleaned up client {client_id}. Remaining connections: {len(active_connections)}")
    except Exception as e:
        logger.error(f"WebSocket error for client {client_id}: {e}", exc_info=True)
        if client_id in active_connections:
            del active_connections[client_id]
        ssh_manager.cleanup_client_connections(client_id)

async def handle_websocket_message(websocket: WebSocket, client_id: str, data: dict):
    """Handle incoming WebSocket messages"""
    message_type = data.get("type")
    logger.debug(f"Handling message type '{message_type}' from client {client_id}")
    
    try:
        if message_type == "ssh_connect":
            await handle_ssh_connect(websocket, client_id, data)
        elif message_type == "terminal_input":
            await handle_terminal_input(websocket, client_id, data)
        elif message_type == "terminal_resize":
            await handle_terminal_resize(websocket, client_id, data)
        elif message_type == "ai_command":
            await handle_ai_command(websocket, client_id, data)
        elif message_type == "ai_chat":
            await handle_ai_chat(websocket, client_id, data)
        elif message_type == "get_suggestions":
            await handle_get_suggestions(websocket, client_id, data)
        elif message_type == "test_connection":
            await handle_test_connection(websocket, client_id, data)
        else:
            logger.warning(f"Unknown message type '{message_type}' from client {client_id}")
    except Exception as e:
        logger.error(f"Error handling message type '{message_type}' from client {client_id}: {e}", exc_info=True)
        await websocket.send_json({
            "type": "error",
            "message": f"Server error: {str(e)}"
        })

async def handle_ssh_connect(websocket: WebSocket, client_id: str, data: dict):
    """Handle SSH connection request"""
    logger.info(f"SSH connection request from client {client_id}")
    try:
        connection_info = ConnectionInfo(**data["connection"])
        logger.info(f"Attempting SSH connection to {connection_info.hostname}:{connection_info.port} for user {connection_info.username}")
        
        # Create SSH connection
        session_id = await ssh_manager.create_connection(
            client_id, connection_info, websocket
        )
        
        # Store connection info in database
        db_manager.store_connection(connection_info)
        logger.info(f"SSH connection successful. Session ID: {session_id}")
        
        await websocket.send_json({
            "type": "ssh_connected",
            "session_id": session_id,
            "success": True
        })
        
    except Exception as e:
        logger.error(f"SSH connection failed for client {client_id}: {e}", exc_info=True)
        await websocket.send_json({
            "type": "ssh_error",
            "error": str(e)
        })

async def handle_terminal_input(websocket: WebSocket, client_id: str, data: dict):
    """Handle terminal input"""
    session_id = data.get("session_id")
    input_data = data.get("input")
    logger.debug(f"Terminal input from client {client_id}, session {session_id}: {repr(input_data)}")
    
    if session_id in ssh_manager.sessions:
        await ssh_manager.send_input(session_id, input_data)
    else:
        logger.warning(f"Session {session_id} not found for client {client_id}")

async def handle_terminal_resize(websocket: WebSocket, client_id: str, data: dict):
    """Handle terminal resize"""
    session_id = data.get("session_id")
    cols = data.get("cols", 80)
    rows = data.get("rows", 24)
    logger.debug(f"Terminal resize from client {client_id}, session {session_id}: {cols}x{rows}")
    
    if session_id and session_id in ssh_manager.sessions:
        await ssh_manager.resize_terminal(session_id, cols, rows)
        logger.info(f"Successfully resized terminal for session {session_id} to {cols}x{rows}")
        # Send confirmation back to client
        await websocket.send_json({
            "type": "terminal_resize_ack",
            "session_id": session_id,
            "cols": cols,
            "rows": rows,
            "success": True
        })
    else:
        logger.debug(f"Terminal resize ignored - no active SSH session for client {client_id} (session_id: {session_id})")

async def handle_ai_command(websocket: WebSocket, client_id: str, data: dict):
    """Handle AI command generation request"""
    try:
        query = data.get("query")
        command_type = data.get("command_type", "linux")
        context = data.get("context", {})
        
        # Generate command using AI
        command = ai_generator.generate_command(query, command_type, context)
        
        # Get related suggestions
        suggestions = ai_generator.get_related_commands(command, command_type)
        
        await websocket.send_json({
            "type": "ai_command_response",
            "query": query,
            "command": command,
            "suggestions": suggestions,
            "command_type": command_type
        })
        
    except Exception as e:
        logger.error(f"AI command generation failed: {e}")
        await websocket.send_json({
            "type": "ai_error",
            "error": str(e)
        })

async def handle_ai_chat(websocket: WebSocket, client_id: str, data: dict):
    """Handle AI chat request with streaming response"""
    try:
        message = data.get("message")
        context = data.get("context", {})
        
        # Send initial response to indicate streaming started
        await websocket.send_json({
            "type": "ai_chat_stream_start",
            "message": message
        })
        
        # Get streaming AI response
        full_response = ""
        for chunk in ai_generator.generate_chat_response_stream(message, context):
            full_response += chunk
            # Send each chunk as it arrives
            await websocket.send_json({
                "type": "ai_chat_stream_chunk",
                "message": message,
                "chunk": chunk,
                "partial_response": full_response
            })
        
        # Send final complete response
        await websocket.send_json({
            "type": "ai_chat_response",
            "message": message,
            "response": full_response
        })
        
    except Exception as e:
        logger.error(f"AI chat failed: {e}")
        await websocket.send_json({
            "type": "ai_error",
            "error": str(e)
        })

async def handle_get_suggestions(websocket: WebSocket, client_id: str, data: dict):
    """Handle request for AI suggestions"""
    try:
        connection_id = data.get("connection_id")
        current_context = data.get("context", {})
        logger.debug(f"Getting suggestions for client {client_id}, connection {connection_id}")
        
        # Get intelligent suggestions based on context
        suggestions = ai_connection_manager.get_contextual_suggestions(
            connection_id, current_context
        )
        
        await websocket.send_json({
            "type": "suggestions_response",
            "suggestions": suggestions
        })
        
    except Exception as e:
        logger.error(f"Get suggestions failed: {e}")
        await websocket.send_json({
            "type": "ai_error",
            "error": str(e)
        })

async def handle_test_connection(websocket: WebSocket, client_id: str, data: dict):
    """Handle test connection message"""
    logger.info(f"Test connection from client {client_id} at {data.get('timestamp')}")
    await websocket.send_json({
        "type": "test_connection",
        "status": "success",
        "server_time": datetime.now().isoformat(),
        "client_id": client_id
    })

# REST API endpoints
@app.post("/api/connections")
async def create_connection(connection: ConnectionInfo):
    """Create new connection"""
    try:
        connection_id = db_manager.store_connection(connection)
        return {"success": True, "connection_id": connection_id}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@app.get("/api/connections")
async def get_connections():
    """Get all connections"""
    try:
        connections = db_manager.get_connections()
        return {"connections": connections}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/enterprises")
async def get_enterprises():
    """Get all enterprises with their connections"""
    try:
        enterprises = ai_connection_manager.get_enterprise_summary()
        return {"enterprises": enterprises}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/ai/command")
async def generate_command(request: CommandRequest):
    """Generate command using AI"""
    try:
        command = ai_generator.generate_command(
            request.query, 
            request.command_type, 
            request.context
        )
        return {"command": command, "type": request.command_type}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/analytics")
async def get_analytics():
    """Get usage analytics"""
    try:
        analytics = db_manager.get_analytics()
        return analytics
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    uvicorn.run(
        "app:app",
        host="0.0.0.0",
        port=8001,
        reload=True,
        log_level="info"
    )